#!/bin/bash

CURRENT_DIR=$(dirname $(realpath "$0"))

################################################

declare -A arch_list=(["mt6877"]="mdla2.0" \
    ["mt6878"]="mdla3.5" \
    ["mt6879"]="mdla2.0" \
    ["mt6897"]="mdla5.0" \
    ["mt8792"]="mdla5.0" \
    ["mt6886"]="mdla3.5" \
    ["mt6893"]="mdla1.5" \
    ["mt6895"]="mdla3.0" \
    ["mt6983"]="mdla3.0" \
    ["mt6985"]="mdla3.5" \
    ["mt6989"]="mdla5.1")
declare -A core_size_list=(["mt6877"]=1 ["mt6878"]=1 ["mt6879"]=1 ["mt6886"]=1 ["mt6893"]=2 ["mt6895"]=2 ["mt6983"]=4 ["mt6985"]=4 ["mt6989"]=4 ["mt6897"]=2 ["mt8792"]=2)
declare -A tcm_size_list=(["mt6877"]=0 ["mt6878"]=0 ["mt6879"]=0 ["mt6886"]=2048 ["mt6893"]=1024 ["mt6895"]=2048 ["mt6983"]=6144 ["mt6985"]=6144 ["mt6989"]=7168 ["mt6897"]=3072 ["mt8792"]=3072)
declare -A out_dir_list=(["mt6877"]="mt6877_g68" \
    ["mt6878"]="mt6878_g615" \
    ["mt6879"]="mt6879_g610" \
    ["mt6886"]="mt6886_g610" \
    ["mt6897"]="mt6897_g615" \
    ["mt8792"]="mt8792_g615" \
    ["mt6893"]="mt6893_g77"  \
    ["mt6895"]="mt6895_g610" \
    ["mt6983"]="mt6983_g710" \
    ["mt6985"]="mt6985_g715_immortalis" \
    ["mt6989"]="mt6989_g720-Immortalis")
declare -A sdk_ver_list=(["5.0.2"]="neuropilot-sdk-basic-5.0.2-build20220217" \
    ["5.0.3"]="neuropilot-sdk-basic-5.0.3-build20220318" \
    ["5.0.4"]="neuropilot-sdk-basic-5.0.4-build20220414" \
    ["6.0.0-beta2"]="neuropilot-sdk-basic-6.0.0-beta2-build20220802" \
    ["6.0.1"]="neuropilot-sdk-basic-6.0.1-build20221013" \
    ["6.0.3"]="neuropilot-sdk-basic-6.0.3-build20230118" \
    ["7.0.2"]="neuropilot-sdk-basic-7.0.2-build20231029" \
	["7.0.8"]="neuropilot-sdk-premium-7.0.8-build20240807" \
    ["8.0.1"]="neuropilot-sdk-premium-8.0.1-build20240703" \
    ["8.0.2"]="neuropilot-sdk-premium-8.0.2-build20240731" \
    ["8.0.4"]="neuropilot-sdk-premium-8.0.4-build20240925")

SDK_PATH="/data2/tools/MTKTools/SDK_Release"
LIB_PATH="neuron_sdk/host/lib"
BIN_PATH="neuron_sdk/host/bin/ncc-tflite"

# @brief 执行转换
function tran_model() {
    local input_file=$1
    local platform=$2
    local enable_tcm=$3
    local en_3=$5
    local en_bw=$6
    local en_sp_in=$7
    local en_sp_out=$8
    local en_debug=$9
    local en_memory=${10}
    local out_dir=$4

    file_ext=${input_file##*.}
    if [[ "${file_ext}" != "tflite" ]];then
        echo "Error: must be tflite model, current is ${file_ext}"
        return -1
    fi

    #file_dir=$(dirname $(realpath $out_dir))
    out_file=${input_file##*/}
    out_file=${out_file%.*}
	

    apu_arch=${arch_list[$platform]}
    core_size=${core_size_list[$platform]}
    tcm_size=${tcm_size_list[$platform]}

    echo "******************************"
    echo "--->begin to trans <${input_file}>, target dir: <${out_dir}>"
    echo "    target platform: ${platform}"
    echo "    arch: ${apu_arch}"
    echo "    core size: ${core_size}"

    extra_cmd=""
    if [[ $enable_tcm == 1 ]];then
        extra_cmd="${extra_cmd} --l1-size-kb ${tcm_size}"
        echo "    tcm open, size: ${tcm_size} KB"
        out_file="${out_file}_tcm"
    fi
    if [[ $en_3 == 1 ]];then
        extra_cmd="${extra_cmd} --opt 3"
        echo "    opt 3 open"
        out_file="${out_file}_opt3"
    fi
    if [[ $en_bw == 1 ]];then
        extra_cmd="${extra_cmd} --opt-bw"
        echo "    opt memory bit width open"
        out_file="${out_file}_opt-bw"
    fi
    if [[ $en_sp_in == 1 ]];then
        extra_cmd="${extra_cmd} --suppress-input"
        echo "    Suppress input data conversion"
        out_file="${out_file}_sp-in"
    fi
    if [[ $en_sp_out == 1 ]];then
        extra_cmd="${extra_cmd} --suppress-output"
        echo "    Suppress output data conversion"
        out_file="${out_file}_sp-out"
    fi
    if [[ $en_debug == 1 ]];then
        extra_cmd="${extra_cmd} --gen-debug-info"
        echo "    Produce debugging information in the DLA file."
        echo "    Runtime can work with this info for profiling."
        out_file="${out_file}_debug"
    fi
    if [[ $en_memory == 1 ]];then
        extra_cmd="${extra_cmd} --show-memory-summary"
        echo "    Show memory allocation summary"
        out_file="${out_file}_memory"
    fi
    out_file=$out_dir/${out_file%.*}.dla
    export $TRAN_PATH
    CMD="$TRAN_CMD --arch=$apu_arch --platform $platform --relax-fp32 $input_file -o $out_file --num-mdla $core_size ${extra_cmd}"
# --int8-to-uint8 --suppress-output"
    echo $CMD
    $CMD

    if [[ $? == 0 ]];then
        echo "--->trans <${input_file}> suc, save to <${out_file}>"
    else
        echo "--->trans <${input_file}>, failed."
    fi
    echo "******************************"
    echo ""

    return 0
}

################################################

# @brief 参数使用说明
function Usage() {
    echo "---------------------------------------------------------------------------------"
    echo "$0"
    echo "Usage:"
    echo "本脚本用于将tflite文件转换为apu目标文件, 架构使用mdla"
    echo "-------------------------------------"
    echo "  --help|-h      查看指令帮助信息"
    echo "-------------------------------------"
    echo "  --dir|-d            [可选], 指定待打包模型的原始目录"
    echo "  --input|-i          [可选], 指定待打包模型文件"
    echo "              以上参数必须指定一个且只能有一个有效"
    echo "  --platform|-p       [必须], 指定待打包模型的目标机型, 例如mt6983"
    echo "  --op3               [可选], 优化模型性能"
    echo "  --opbw              [可选], 优化转换模型的内存位宽"
    echo "  --sdk|-v            [可选], 默认版本为6.0.3"
    echo "  --tcm|-t            [可选], 指定开启TCM, 默认关闭"
    echo "  --suppress-input    [可选], 抑制输入数据转换,float模型配合--relax-fp32开启后输入数据为fp16"
    echo "  --suppress-output   [可选], 抑制输出数据转换,float模型配合--relax-fp32开启后输出数据为fp16,仅NP6及以上版本会生效"
    echo "---------------------------------------------------------------------------------"
}

################################################

input_dir=""
input_file=""
tar_plat=""
sdk_version="6.0.3"
enable_tcm=0
enable_op3=0
enable_opbw=0
enable_suppress_input=0
enable_suppress_output=0
enable_debug=0
enable_memory=0

while [ -n "$1" ]
do
  case "$1" in
    --dir|-d)
        input_dir=$2
        if [[ ! -d "${input_dir}" ]];then
            echo "Error: input dir: <${input_dir}> is not a valid dir"
            Usage
            exit 1
        fi
        shift
        ;;
    --input|-i)
        input_file=$2
        if [[ ! -f "${input_file}" ]];then
            echo "Error: input file: <${input_dir}> is not exist"
            Usage
            exit 1
        fi
        input_dir=$(dirname $input_file)
        shift
        ;;
    --platform|-p)
        tar_plat=$2
        shift
        ;;
    --op3)
        enable_op3=1
        ;;
    --opbw)
        enable_opbw=1
        ;;
    --sdk|-v)
        sdk_version=$2
        shift
        ;;
    --tcm|-t)
        enable_tcm=1
        ;;
    --suppress-input)
        enable_suppress_input=1
        ;;
    --suppress-output)
        enable_suppress_output=1
        ;;
    --debug)
        enable_debug=1
        ;;
    --memory)
        enable_memory=1
        ;;
    --help|-h)
        Usage
        exit 0
        ;;
    *)
        echo "Error: unsupport:$1"
        Usage
        exit 1
        ;;
    esac
    shift
done

################################################

echo "input file: <${input_file}>"
echo "input dir: <${input_dir}>"
echo "platform: <${tar_plat}>"
echo "sdk version: <${sdk_version}>"
echo "enable_tcm: <${enable_tcm}>"
echo "enable_op3: <${enable_op3}>"
echo "enable_opbw: <${enable_opbw}>"
echo "enable_suppress_input: <${enable_suppress_input}>"
echo "enable_suppress_output: <${enable_suppress_output}>"
echo "enable_debug: <${enable_debug}>"
echo "enable_memory: <${enable_memory}>"

if [[ -z "${input_dir}" ]] && [[ -z "${input_file}" ]];then
    echo "Error: --input and --dir must be set one"
    Usage
    exit 1
fi

has_plat=0
for plat in ${!arch_list[@]}
do
    if [[ $plat == $tar_plat ]];then
        has_plat=1
        break;
    fi
done
if [[ $has_plat == 0 ]];then
    echo "Error: target platform <${tar_plat}> invalid"
    exit 1
else
    if ([[ "$tar_plat" == "mt6985" ]] || [[ "$tar_plat" == "mt6886" ]]) && [[ "$sdk_version" =~ "5.0." ]];then
        echo "Error: sdk version need >= 6.0.0 for paltform mt6985 or mt6886"
        exit 1
    fi
fi
if [[ $enable_suppress_output == 1 ]] && [[ "$sdk_version" =~ "5.0." ]];then
    echo "Error: neuropilot-sdk need >= 6.0.0 for --suppress-output"
    # enable_suppress_output=0
    exit 1
fi
################################################

TRAN_PATH="LD_LIBRARY_PATH=${SDK_PATH}/${sdk_ver_list[$sdk_version]}/${LIB_PATH}"
TRAN_CMD="${SDK_PATH}/${sdk_ver_list[$sdk_version]}/${BIN_PATH}"

current=`date +%Y-%m-%d-%H-%M-%S`
out_name=${out_dir_list[$tar_plat]}_${current}
out_dir=$input_dir/${out_name}
mkdir -p ${out_dir}
echo "====>mkdir ${out_dir}"
ret=0
if [[ -f "${input_file}" ]];then
    tran_model ${input_file} ${tar_plat} ${enable_tcm} ${out_dir} ${enable_op3} ${enable_opbw} ${enable_suppress_input} ${enable_suppress_output} ${enable_debug} ${enable_memory}
    ret=$?
elif [[ -d "${input_dir}" ]];then
    for tfile in ${input_dir}/*.tflite
    do
        tran_model ${tfile} ${tar_plat} ${enable_tcm} ${out_dir} ${enable_op3} ${enable_opbw} ${enable_suppress_input} ${enable_suppress_output} ${enable_debug} ${enable_memory}
        ret=$?
        if [[ $ret -ne 0 ]];then
            break;
        fi
    done
else
    echo "Error: could not found valid input tflite file"
    ret=1
fi

if [[ $ret -ne 0 ]];then
    echo "Error: tran model FAILED!"
else
    cd $input_dir
    out_zip="${out_name}.zip"
    out_md5="${out_name}.md5"
    echo "zip ${out_zip} ${out_name}"
    zip -r ${out_zip} ${out_name}
    md5sum ${out_zip} | cut -d ' ' -f 1 > ${out_md5}
    echo "md5:$(cat ${out_md5})"
    echo "mv ${out_zip} ${CURRENT_DIR}/.output_archive/"
    echo "mv ${out_md5} ${CURRENT_DIR}/.output_archive/"
    mv ${out_zip} ${CURRENT_DIR}/.output_archive/
    mv ${out_md5} ${CURRENT_DIR}/.output_archive/
fi